import torch
import torch.nn as nn
import torch.optim as optim
import os
import math
import random
import argparse
import numpy as np
import rospy
from std_msgs.msg import String
from std_msgs.msg import Header
from sensor_msgs.msg import PointCloud2
import sensor_msgs.point_cloud2 as pc2
from utils import utils_ros
from model import CarSpeedNet,FrogEyeNet,FrogEyeDetNet

def ForgEye_val(model,x_batch):
    model.eval()
    # 测试模型
    print(x_batch.shape)
    x_batch = torch.from_numpy(x_batch).unsqueeze(0)
    print(x_batch.shape)
    x_batch = x_batch.to(torch.float32).to("cuda:0")
    print(x_batch.shape)
    output, xv_car = model(x_batch)
    predict = nn.functional.softmax(output, dim=2) #转换为概率
    predict = torch.argmax(predict, dim=2) #取概率高状态为实际状态
    print(predict.shape)
    return predict.squeeze().detach().cpu().numpy()

pub = rospy.Publisher('radar_raw', PointCloud2, queue_size=5)
pub_out = rospy.Publisher('radar_out', PointCloud2, queue_size=5)
# 初始化模型
carspeed_model = CarSpeedNet()
mymodel = FrogEyeNet(pre_model=None)

def ForgEye(config, model):
    device = config.device
    print("using {} device.".format(device))
    model.to(device)
    # 读取模型的预训练模型
    model_path = os.path.join(args.checkpoints_dir,
                              "Model_" + "FrogEye" + "_epoch_" + str(args.resume_epoch) + ".pth")
    assert os.path.isfile(model_path), "model_path <{}> does not exist.".format(model_path)
    model.load_state_dict(torch.load(model_path, map_location=args.device))

def callback(data):
    pc = pc2.read_points(data, skip_nans=True, field_names=("x", "y", "z", "normal_x","normal_y","normal_z","intensity","curvature"))
    pc_list = []
    for p in pc:
        pc_list.append([p[0], p[1], p[2], 0,p[3],p[4], p[5],0, p[6], p[7],0,0])#注意强度归一化
    lidarData = np.array(pc_list)
    # print("data:{}".format(lidarData[0,:]))

    x = lidarData[:,(0,1,2,4)]
    # 归一化处理
    x[0] = x[0] / 100.0
    x[1] = x[1] / 100.0
    x[2] = (x[2] + 1.0) / 2.0
    x[3] = (x[3] + 30) / 30.0
    # print("x:{}".format(x[0,:]))
    x = x[:100,:]
    print("ssss:{}".format(x.shape))
    pred = ForgEye_val(mymodel,x)
    print("pred:{}".format(pred.shape))
    pub.publish(utils_ros.array2msg(lidarData))#发布原始点云
    output = lidarData[:100]
    output[:,8] = pred
    print("output[:,8]:{}".format(output[:,8].shape))
    pub_out.publish(utils_ros.array2msg(output))#发布预测点云


def listener():
    rospy.init_node('listener', anonymous=True)
    rospy.Subscriber("/radar_pc_diff", PointCloud2, callback)
    rospy.spin()



if __name__ == '__main__':

    # 相对根路径
    root_path = "/home/zwh/Desktop/testAndlearn/learn/DeepLearn学习/deep-learning-for-image-processing"  # get data root path   os.getcwd()为执行终端路径
    assert os.path.exists(root_path), "{} path does not exist.".format(root_path)
    parser = argparse.ArgumentParser()
    parser.add_argument('--root-path', type=str, default=root_path)
    parser.add_argument('--current-path', type=str, default=root_path + "/pytorch_classification/FrogEye")
    parser.add_argument('--checkpoints_dir', type=str, default=root_path + "/pytorch_classification/FrogEye/weights")
    parser.add_argument('--resume_epoch', type=int, default=100000)
    parser.add_argument('--device', type=str, default='cuda:0')
    parser.add_argument('--batch-size', type=int, default=2)
    parser.add_argument('--num_work', type=int, default=1)
    parser.add_argument('--lr', type=float, default=0.001)
    parser.add_argument('--lrf', type=float, default=0.01)
    parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
                        help='momentum')
    parser.add_argument('--wd', '--weight-decay', default=1e-4, type=float,
                        metavar='W', help='weight decay (default: 1e-4)',
                        dest='weight_decay')
    parser.add_argument('--start_epoch', type=int, default=0)
    parser.add_argument('--end_epoch', type=int, default=500)
    parser.add_argument('--prune_epoch', type=int, default=50)
    parser.add_argument('--tb_writer', type=bool, default=True)
    parser.add_argument('--val_start_epoch', type=int, default=20)
    parser.add_argument('--checkpoint_freq', type=int, default=10)
    parser.add_argument('--FrogEye_Wights', type=bool, default=True)
    args = parser.parse_args()

    ForgEye(args,mymodel)

    listener()


